import os

import mmcv
# using a pre-trained detector. 预训练配置
from mmcv import Config


# root_path=r"/project/train/src_repo/mmdetection/" # 要改的

root_path=r"/home/deepin/Documents/ji_pingtai/mmdetection/" # 要改的

cfg_path=root_path+'configs/swin/faster_rcnn_swin_t-p4-w7_fpn_3x_coco.py'

cfg = Config.fromfile(cfg_path) # 加载cfg  文件

from mmdet.apis import set_random_seed, inference_detector, show_result_pyplot

# # Modify dataset type and path
# cfg.dataset_type = 'CocoDataset' # 自定义的数据集的 类
# cfg.data_root = root_path+'data/coco/'

# # 训练集
# cfg.data.train.type = 'CocoDataset'
# cfg.data.train.data_root = cfg.data_root 
# cfg.data.train.ann_file = 'annotations/instances_train2017.json' # 存 图片名字
# cfg.data.train.img_prefix = 'train2017'# 图片位置

# #验证集
# cfg.data.val.type = 'CocoDataset'
# cfg.data.val.data_root = cfg.data_root 
# cfg.data.val.ann_file = 'annotations/instances_val2017.json'
# cfg.data.val.img_prefix = 'val2017'

# # 测试集
# cfg.data.test.type = 'CocoDataset'
# cfg.data.test.data_root =cfg.data_root 
# cfg.data.val.ann_file = 'annotations/instances_val2017.json'
# cfg.data.val.img_prefix = 'val2017'



cfg.device='cuda' # 'ConfigDict' object has no attribute 'device'

# # modify num classes of the model in box head 类别数
# cfg.model.roi_head.bbox_head.num_classes = 2 #----类别数

# # If we need to finetune a model based on a pre-trained detector, we need to
# # use load_from to set the path of checkpoints. 预训练模型 
# #下载： wget http://download.openmmlab.com/mmdetection/v2.0/faster_rcnn/faster_rcnn_r50_fpn_1x_coco/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth

# cfg.load_from = root_path+'checkpoints/3x_coco.pth'


# # Set up working dir to save files and logs. 保存的文件夹
cfg.work_dir = 'work_dir'

# # The original learning rate (LR) is set for 8-GPU training.
# # We divide it by 8 since we only use one GPU.
# #原始学习率（LR）被设置用于8-GPU训练。
# #我们将其除以8，因为我们只使用一个GPU。
# cfg.optimizer.lr = 0.05 / 8 # 过大，学习率过高，梯度下降快，梯度消失
# cfg.lr_config.warmup = None
cfg.log_config.interval = 20 #打印log 间隔

# # Change the evaluation metric since we use customized dataset.
# #更改评估指标，因为我们使用自定义数据集。
# # cfg.evaluation.metric = 'mAP'

# # We can set the evaluation interval to reduce the evaluation times

cfg.evaluation.interval = 10 # 评估间隔，以减少评估时间

# # We can set the checkpoint saving interval to reduce the storage cost

cfg.checkpoint_config.interval = 25  # 保存间隔:  设置检查点：，以降低存储成本

cfg.runner = dict(type='EpochBasedRunner', max_epochs=30)# ---训练 轮次


# # Set seed thus the results are more reproducible #播种，结果更具可重复性
cfg.seed = 0
set_random_seed(0, deterministic=False)
cfg.gpu_ids = range(1)

# # We can also use tensorboard to log the training process #我们还可以使用tensorboard记录培训过程
# cfg.log_config.hooks = [
#     dict(type='TextLoggerHook'),
#     dict(type='TensorboardLoggerHook')]


# We can initialize the logger for training and have a look
# at the final config used for training
#我们可以初始化记录器进行培训并查看
#在用于培训的最终配置中
print(f'Config:\n{cfg.pretty_text}')



#训练----------------------------------------------------
from mmdet.datasets import build_dataset
from mmdet.models import build_detector
from mmdet.apis import train_detector


# Build dataset 构建 数据集
datasets = [build_dataset(cfg.data.train)]

# Build the detector 构建 识别
model = build_detector(cfg.model)

# Add an attribute for visualization convenience 添加属性以方便可视化，模型的类别
model.CLASSES = datasets[0].CLASSES

# Create work_dir 文件夹
mmcv.mkdir_or_exist(os.path.abspath(cfg.work_dir))

#保存 cfg 配置：
# 把自定义的 config 在 自定义的working dir 下 保存一份------------推理和测试 可以用到的
cfg.dump(F'{cfg.work_dir}/swin_cfgformat.py')

train_detector(model, datasets, cfg, distributed=False, validate=True) # -------进行----训练-----


# # # 查看 评估结果 load tensorboard in colab
# # %load_ext tensorboard
# #
# # # see curves in tensorboard
# # tensorboard --logdir ./tutorial_exps  #--------运行这个看看 数的 曲线



